package com.mini.grpc.transport.netty;

import com.mini.grpc.common.Metadata;
import com.mini.grpc.common.MethodDescriptor;
import com.mini.grpc.common.Status;
import com.mini.grpc.transport.ClientStream;
import com.mini.grpc.transport.ClientStreamListener;
import io.netty.handler.codec.http2.Http2Headers;
import io.netty.handler.codec.http2.DefaultHttp2Headers;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Netty客户端流实现
 */
public class NettyClientStream<ReqT, RespT> implements ClientStream<ReqT, RespT> {
    
    private static final Logger logger = LoggerFactory.getLogger(NettyClientStream.class);
    
    private final int streamId;
    private final NettyClientHandler handler;
    private final MethodDescriptor<ReqT, RespT> method;
    private ClientStreamListener<RespT> listener;
    private boolean headersSent = false;
    private boolean closed = false;
    
    public NettyClientStream(int streamId, NettyClientHandler handler, MethodDescriptor<ReqT, RespT> method) {
        this.streamId = streamId;
        this.handler = handler;
        this.method = method;
    }
    
    public void setListener(ClientStreamListener<RespT> listener) {
        this.listener = listener;
    }
    
    @Override
    public void sendMessage(ReqT message) {
        if (closed) {
            throw new IllegalStateException("Stream is closed");
        }
        
        // 如果还没有发送头部，先发送默认头部
        if (!headersSent) {
            sendHeaders(new Metadata());
        }
        
        try {
            // 序列化消息并按照 gRPC 格式发送
            // gRPC 消息格式：压缩标志(1字节) + 消息长度(4字节) + 消息内容
            byte[] messageBytes = message.toString().getBytes("UTF-8");
            
            // 创建 gRPC 消息帧
            byte[] grpcMessage = new byte[5 + messageBytes.length];
            
            // 压缩标志（1字节，0表示不压缩）
            grpcMessage[0] = 0;
            
            // 消息长度（4字节，大端序）
            int length = messageBytes.length;
            grpcMessage[1] = (byte) ((length >> 24) & 0xFF);
            grpcMessage[2] = (byte) ((length >> 16) & 0xFF);
            grpcMessage[3] = (byte) ((length >> 8) & 0xFF);
            grpcMessage[4] = (byte) (length & 0xFF);
            
            // 消息内容
            System.arraycopy(messageBytes, 0, grpcMessage, 5, messageBytes.length);
            
            logger.debug("Sending gRPC message for stream {}: {} bytes", streamId, grpcMessage.length);
            handler.sendData(streamId, grpcMessage, false);
        } catch (Exception e) {
            logger.error("Failed to send message", e);
            if (listener != null) {
                listener.closed(Status.INTERNAL.withCause(e), new Metadata());
            }
        }
    }
    
    @Override
    public void request(int numMessages) {
        // 客户端流不需要实现这个方法
    }
    
    @Override
    public void halfClose() {
        if (closed) {
            return;
        }
        
        try {
            handler.sendData(streamId, new byte[0], true);
        } catch (Exception e) {
            logger.error("Failed to half close stream", e);
        }
    }
    
    @Override
    public void cancel(Status status, Metadata trailers) {
        if (closed) {
            return;
        }
        
        closed = true;
        try {
            handler.cancelStream(streamId);
            if (listener != null) {
                listener.closed(status, trailers);
            }
        } catch (Exception e) {
            logger.error("Failed to cancel stream", e);
        }
    }
    
    @Override
    public int getStreamId() {
        return streamId;
    }
    
    public void sendHeaders(Metadata headers) {
        logger.info("=== NettyClientStream.sendHeaders called for stream {} ===", streamId);
        
        if (headersSent) {
            logger.error("Headers already sent for stream {}", streamId);
            throw new IllegalStateException("Headers already sent");
        }
        
        Http2Headers http2Headers = convertMetadata(headers);
        
        // 添加必要的 gRPC HTTP/2 头部
        http2Headers.method("POST");
        http2Headers.path(method.getFullMethodName());
        http2Headers.scheme("http");
        http2Headers.add("content-type", "application/grpc");
        http2Headers.add("te", "trailers");
        
        logger.info("Sending headers for stream {}: {}", streamId, http2Headers);
        
        try {
            handler.sendHeaders(streamId, http2Headers, false);
            headersSent = true;
            logger.info("Headers sent successfully for stream {}", streamId);
        } catch (Exception e) {
            logger.error("Failed to write headers for stream {}", streamId, e);
            if (listener != null) {
                listener.closed(Status.INTERNAL.withCause(e), new Metadata());
            }
        }
    }
    
    // 由NettyClientHandler调用的方法
    public void onHeaders(Metadata headers) {
        if (listener != null) {
            listener.headersReceived(headers);
        }
    }
    
    public void onMessage(RespT message) {
        if (listener != null) {
            listener.messageReceived(message);
        }
    }
    
    public void onComplete() {
        if (listener != null) {
            listener.closed(Status.OK, new Metadata());
        }
        closed = true;
        handler.removeStream(streamId);
    }
    
    public void onError(Status status, Metadata trailers) {
        if (listener != null) {
            listener.closed(status, trailers);
        }
        closed = true;
        handler.removeStream(streamId);
    }
    
    private Http2Headers convertMetadata(Metadata metadata) {
        Http2Headers headers = new DefaultHttp2Headers();
        
        for (String key : metadata.keys()) {
            if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
                // 二进制头
                byte[] value = metadata.get(Metadata.Key.ofBinary(key));
                if (value != null) {
                    headers.add(key, java.util.Base64.getEncoder().encodeToString(value));
                }
            } else {
                // ASCII头
                String value = metadata.get(Metadata.Key.of(key));
                if (value != null) {
                    headers.add(key, value);
                }
            }
        }
        
        return headers;
    }
} 